package yago;
import java.io.File;
import java.io.IOException;
import java.io.Reader;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.StringReader;
import java.util.Collections;
import java.util.Set;
import java.util.HashSet;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.ArrayList;

import org.apache.lucene.util.Version;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Fieldable;
import org.apache.lucene.document.Field.Index;
import org.apache.lucene.document.Field.Store;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.NumericField;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.store.MMapDirectory;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermDocs;
import org.apache.lucene.index.TermFreqVector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.FieldCache;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.QueryTermVector;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.ngram.NGramTokenizer;
import org.apache.lucene.analysis.TokenStream;
import org.xml.sax.Attributes;
import org.xml.sax.SAXException;
import org.xml.sax.helpers.DefaultHandler;




//add by lei yao
import java.io.FileWriter;
import java.io.BufferedWriter;
//////////////////////

import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.SAXParser;
import javax.xml.parsers.SAXParserFactory;

public class IndexTACWriter {
	private static final int NGRAM = 2;

	private IndexWriter writer = null;
	

	private Analyzer analyzer = null;
	private String indexDir;
	private String fileDir;

	private class NGramAnalyzer extends Analyzer {
		private int ngram = 0;
		
		public NGramAnalyzer(int ngram) {
			this.ngram = ngram;
		}
		
		public TokenStream tokenStream(String fieldName, Reader reader) {
			return new NGramTokenizer(reader, 2, ngram);
		}
	}
	
	public IndexTACWriter(String indexLoc,String dir) {
		analyzer = new NGramAnalyzer(NGRAM);
		this.indexDir = indexLoc;
		this.fileDir = dir;
	}

	/**
	 * Init the index writer for the indexing. 
	 * 
	 * @param dirLoc The directory for storing the lucene index.
	 */
	public void initWriter(String dirLoc) {
		Directory dir = null;
		IndexWriterConfig conf = new IndexWriterConfig(Version.LUCENE_34, analyzer); 
		
		
		try {
			dir = new MMapDirectory(new File(dirLoc));
			writer = new IndexWriter(dir, conf);
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * Close and finalize the index writer.
	 */
	public void closeWriter() {
		try {
			writer.optimize();
			writer.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	
	public void indexOnefile(String path){
		
		try {
			SAXParserFactory factory = SAXParserFactory.newInstance();
			SAXParser saxParser = factory.newSAXParser();
			IndexTACHandler handler = new IndexTACHandler(this);
			try {
				saxParser.parse(path, handler);
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
			
		} catch (ParserConfigurationException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (SAXException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
	
	public void addDocument(String wikititle, String content){
		try {
			String name = wikititle;
			
			content = content.replaceAll("\\n", " ");
			Document doc = new Document();
			Fieldable field = null;

			field = new Field("docID", name, Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS);
			doc.add(field);
			field = new Field("content", content, Field.Store.YES, Field.Index.NO);
			doc.add(field);

			writer.addDocument(doc);
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public static File[] getAllFiles(String dir){
		File directory = new File(dir);
        File[] files = directory.listFiles();
    //    System.out.println(files.length);
        for (int i = 0; i < files.length; i++) {
        //	System.out.println(i + ":" + files[i].getName() + ";;;"+ files[i].getAbsolutePath());
           //files[i].renameTo(new File(dir + i + files[i].getName()));
   
           
           //System.out.println(files[i].getName());
       }
       return files;
}
	
	public void write(){
		this.initWriter(this.indexDir);
		
		File[] files = IndexTACWriter.getAllFiles(this.fileDir);
		
		for(int i=0; i<files.length;i++){
			
			System.out.println(i+"/"+files.length);
			String abspath = files[i].getAbsolutePath();
			String path = files[i].getName();
			try {
				this.indexOnefile(abspath);
				
			} catch (Exception e) {
				// TODO Auto-generated catch block
				System.out.println("can not read file:"+abspath);
				e.printStackTrace();
			}
		}
    	
		
		
		this.closeWriter();
	}
	


	
	public static void main(String[] args) {
		IndexTACWriter writer = new IndexTACWriter("/largedata1/cmput696/students/lyao1/tacindex","/largedata1/cmput696/knowledge_base/tac-kbp-knowledge-base-2011/data");
		writer.write();
		System.out.println("written finished"); 
	}
}

